@DejunL thanks for the contribution, but this PR is too large to review. It needs to be broken up into more readily reviewable chunks. If you want more context, please refer to the contributing and PR review guides in the repository. Please aim for 500-1000 lines at the maximum to make this tractable for reviewers. Thanks!

After a first pass, my suggestion is to do this by:

* having a PR for the webdataset

* having another PR for diffdock utils

Webdataset looks like it is diffdock specific, so please put this into bionemo-diffdock and hold off on putting it into bionemo-core until we have confirmed multiple use cases.

Additionally, the diffdock utils need unit tests before they can be merged in.

Thanks for the suggestion. I should clarify the motivation in the OP but let me do it here too. The webdatamodule is not diffdock specific. Nothing in the datamodule.py depends on any specific model. DiffDock-pp, which we are planning to add to BioNeMo soon, is also using webdataset. @kdidiNVIDIA is also interested in trying to use webdataset in the protein foundation model. So putting webdatamodule into bionemo-diffdock will not achieve the desirable modularization.

My original thinking of putting webdatamodule into bionemo-core was seeing the PRNGDatasetShuffler in it. The webdataset data processing routines share the same or even higher level of generalizability (considering PRNGDatasetShuffler only works for indexable dataset but not iterable dataset while webdataset can handle both). So I think bionemo-core would be a good place to start with.

In the context of @guoqing-zhou working on modularizing DiffDock specific code, I can remove the the diffdock tests and utils from this PR and only focus on the generic webdatamodule. This would also reduce the number of lines to review. What do you think? @malcolmgreaves

#100 (comment)

Following up from the slack discussion - let's double check alignment and summarize next steps:
(1) Make a new sub-package for the webdataset dataloader. Make a PR with this.
(2) Make a 2nd PR for Diffdock. It will depend on this webdataset dataloader subpackage.

How does this sound @DejunL ?

Following up from the slack discussion - let's double check alignment and summarize next steps:
(1) Make a new sub-package for the webdataset dataloader. Make a PR with this.
(2) Make a 2nd PR for Diffdock. It will depend on this webdataset dataloader subpackage.

First of all, we should stop referring this PR to "dataset" or "dataloader" as this is about lightning data module for using webdataset.

I have just moved the code under bionemo-webdatamodule and remove the diffdock specific code because the webdatamodule doesn't depend on diffdock at all (the original code you were referring to was only used for testing webdatamodule). The diffdock migration to bionemo2 will be future work. @malcolmgreaves

/build-ci